# 实现网格绘制
# 注意到流体网格是完全不变的 只需要研究固体网格

import dolfinx
from mpi4py import MPI
import numpy as np
import gmsh
import os

# 切换工作目录到当前脚本所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(current_dir)

n_fluid = 16  # 流体网格单元数
n_solid1 = 8   # 固体网格1单元数
n_solid2 = 4   # 固体网格2单元数 
pi = np.pi
# 流体网格：[-2,2]^2，16x16三角形网格
fluid_mesh = dolfinx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-2, -2]), np.array([2, 2])],
    [n_fluid, n_fluid],
    dolfinx.mesh.CellType.triangle
)





# 固体网格：[-1,1]^2，8x8三角形网格（更精细）
solid_mesh1 = dolfinx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-1, -1]), np.array([1, 1])],
    [n_solid1, n_solid1],
    dolfinx.mesh.CellType.triangle
)
solid_mesh2 = dolfinx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-1, -1]), np.array([1, 1])],
    [n_solid2, n_solid2],
    dolfinx.mesh.CellType.triangle,
    diagonal=dolfinx.mesh.DiagonalType.left,
)
solid_mesh3 = dolfinx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-5/8, -5/8]), np.array([13/8, 13/8])],
    [n_solid2, n_solid2],
    dolfinx.mesh.CellType.triangle,
    diagonal=dolfinx.mesh.DiagonalType.left,
)
# 读取gmsh 正方形网格
gmsh.initialize()
gmsh.open("square_unstr_solid.msh")  
solid_mesh4, cell_marker, facet_marker = dolfinx.io.gmshio.model_to_mesh(
    gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
gmsh.finalize()
solid_mesh6 = dolfinx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-pi/4, -pi/4]), np.array([pi/4, pi/4])],
    [n_solid2, n_solid2],
    dolfinx.mesh.CellType.triangle,
    diagonal=dolfinx.mesh.DiagonalType.left,
)
# 读取gmsh 圆盘网格
gmsh.initialize()
gmsh.open("disk_solid.msh")  # 确保当前目录下有disk.msh文件
solid_mesh7, cell_marker, facet_marker = dolfinx.io.gmshio.model_to_mesh(
    gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
gmsh.finalize()

solid_mesh8 = dolfinx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([-0.62, -0.62]), np.array([1.38, 1.38])],
    [n_solid2, n_solid2],
    dolfinx.mesh.CellType.triangle,
    diagonal=dolfinx.mesh.DiagonalType.left,
)


# 导出为XDMF
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "fluid.xdmf", "w") as xdmf:
    xdmf.write_mesh(fluid_mesh)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid1.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh1)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid2.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh2)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid3.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh3)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid4.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh4)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid6.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh6)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid7.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh7)
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solid8.xdmf", "w") as xdmf:
    xdmf.write_mesh(solid_mesh8)